from common.clickhouse_alc_engine import (
    get_db,
    get_db_engine,
    get_factor_tb,
    get_feature_tb,
    get_kline_tb,
    get_quant_db,
)
from common.tools import limit_iso_str_to_second_str
from common.variables import SymbolType, time_counts, TimeFrameType
from datetime import datetime, timedelta
import pandas as pd
from sqlalchemy import text


def get_symbol_type(symbol: str) -> SymbolType:
    if "/" in symbol:
        return "coin"
    return "stock"


def get_kline_unfilled_range(
    start: str, end: str, symbol: str, timeframe: TimeFrameType
) -> pd.DataFrame:
    db: str = get_db()
    table: str = get_kline_tb(timeframe)
    start = limit_iso_str_to_second_str(start)
    end = limit_iso_str_to_second_str(end)
    sql_query = f"""
    with has_date as (select toDateTime(time) as time, row_number() OVER (ORDER BY time) as idx
                  from {table}
                  where symbol = :symbol
                    and time >= :start
                    and time < :end),

     deltas as (select a.time as start, b.time as end, dateDiff('minute', a.time, b.time) as delta
                from has_date a
                         join has_date b on toInt64(a.idx) = b.idx - 1
                order by start),
     ext_deltas as (select toDateTime(:start)      as start,
                           (select any(deltas.start) from deltas) as end,
                           dateDiff('minute', start, end)         as delta
                    union all
                    select *
                    from deltas
                    union all
                    select (select anyLast(deltas.end) from deltas) as start,
                           toDateTime(:end)        as end,
                           dateDiff('minute', start, end)           as delta),

     check_ranges as (select dateAdd(greatest(start, toDateTime(:start)), interval :minute_cnt minute)as start,
                             end,
                             delta
                      from ext_deltas
                      where delta > :minute_cnt)
    select start, end, delta
    from check_ranges order by start;
    """

    engine = get_db_engine(db)
    with engine.connect() as connection:
        query = connection.execute(
            text(sql_query),
            {
                "start": start,
                "end": end,
                "minute_cnt": time_counts[timeframe] // 60,
                "symbol": symbol,
            },
        )
        results = pd.DataFrame(query.fetchall())
    return results


def get_factor_unfilled_range(
    start: str,
    end: str,
    symbol: str,
    fname_list: list[str],
    timeframe: TimeFrameType,
) -> pd.DataFrame:
    if get_symbol_type(symbol[0]) == "coin":
        db: str = get_db()
        table: str = get_factor_tb(timeframe)
    else:
        db: str = get_quant_db()
        table: str = get_feature_tb(timeframe)
    start = limit_iso_str_to_second_str(start)
    end = limit_iso_str_to_second_str(end)
    sql_query = f"""
    with has_date0 as (select distinct toDateTime(time) as time
                  from {table}
                  where symbol = :symbol
                    and time >= :start
                    and time < :end
                    and fname in :fname_list
                 group by time
                 having count(distinct symbol, fname) >= :fname_list_len

                    ),
    has_date as (
        select time, row_number() OVER (ORDER BY time) as idx
        from has_date0
    ),

     deltas as (select a.time as start, b.time as end, dateDiff('minute', a.time, b.time) as delta
                from has_date a
                         join has_date b on toInt64(a.idx) = b.idx - 1
                order by start),
     ext_deltas as (select toDateTime(:start)      as start,
                           (select any(deltas.start) from deltas) as end,
                           dateDiff('minute', start, end)         as delta
                    union all
                    select *
                    from deltas
                    union all
                    select (select anyLast(deltas.end) from deltas) as start,
                           toDateTime(:end)        as end,
                           dateDiff('minute', start, end)           as delta),

     check_ranges as (select dateAdd(greatest(start, toDateTime(:start)), interval :minute_cnt minute)as start,
                             end,
                             delta
                      from ext_deltas
                      where delta > :minute_cnt)
    select start, end, delta
    from check_ranges order by start;
    """

    engine = get_db_engine(db)
    with engine.connect() as connection:
        query = connection.execute(
            text(sql_query),
            {
                "start": start,
                "end": end,
                "fname_list": fname_list,
                "fname_list_len": len(fname_list),
                "minute_cnt": time_counts[timeframe] // 60,
                "symbol": symbol,
            },
        )
        results = pd.DataFrame(query.fetchall())
    return results


def cut_date_range(
    start: str, end: str, timeframe: TimeFrameType, N: int
) -> list[tuple[datetime, datetime]]:
    """
    :param N: 每个区间段的数量
    """
    s = datetime.fromisoformat(start)
    e = datetime.fromisoformat(end)
    delta_min = time_counts[timeframe] // 60 * N
    p = s
    results = []
    while p < e:
        next_p = p + timedelta(minutes=delta_min)
        if next_p <= e:
            results.append((p, next_p))
        else:
            results.append((p, e))
        p = next_p
    return results


def get_cut_range_from_table(
    df: pd.DataFrame, timeframe: TimeFrameType, N: int
) -> list[tuple[datetime, datetime]]:
    results = []
    for i, row in df.iterrows():
        sp = cut_date_range(row["start"], row["end"], timeframe, N)
        results.extend(sp)
    return results


def fill_data(df: pd.DataFrame, start: str, end: str, timeframe: TimeFrameType):
    results = df
    results.set_index("time", inplace=True)
    real_start = datetime.fromtimestamp(
        ((datetime.fromisoformat(start).timestamp() - 1) // time_counts[timeframe] + 1)
        * time_counts[timeframe]
    )

    need_idx = pd.date_range(real_start, end, freq=f"{time_counts[timeframe]}S")
    need_idx.name = "time"
    results = results.reindex(need_idx)
    # results.interpolate(method='linear', inplace=True)
    results = results.ffill().bfill()
    results.reset_index(inplace=True)
    return results


def write_to_feature_tb(symbol: str, df: pd.DataFrame, timeframe: TimeFrameType):
    df = df.melt(id_vars="time", var_name="fname", value_name="value")
    df["symbol"] = symbol
    if get_symbol_type(symbol) == "stock":
        engine = get_db_engine(get_quant_db())
        df.to_sql(get_feature_tb(timeframe), engine, index=False, if_exists="append")
    else:
        engine = get_db_engine(get_db())
        df.to_sql(get_factor_tb(timeframe), engine, index=False, if_exists="append")


def get_klines(
    start: str,
    end: str,
    symbol: str | list[str],
    timeframe: TimeFrameType,
    is_fill: bool = False,
) -> dict[str, pd.DataFrame]:
    if isinstance(symbol, str):
        symbol = (symbol,)
    db: str = get_db()
    table: str = get_kline_tb(timeframe)
    sql_query = f"""
        select symbol, time, anyLast(open) as open, anyLast(high) as high, anyLast(low) as low,
            anyLast(close) as close, anyLast(volume) as volume from {table}
        where time>=:start and time<:end and symbol in :symbol group by symbol, time order by symbol, time
    """
    engine = get_db_engine(db)
    with engine.connect() as connection:
        query = connection.execute(
            text(sql_query), {"start": start, "end": end, "symbol": symbol}
        )
        results = pd.DataFrame(query.fetchall())
    results = {x: y.drop("symbol") for x, y in results.groupby("symbol")}
    if is_fill:
        results = {x: fill_data(y, start, end, timeframe) for x, y in results}

    return results


def get_factors(
    start: str,
    end: str,
    symbol: str | list[str],
    fname: str | list[str],
    timeframe: TimeFrameType,
    is_fill: bool = False,
) -> dict[str, pd.DataFrame]:
    # 注意：只能一次性输入同一个类型的symbol
    if isinstance(symbol, str):
        symbol = (symbol,)
    if isinstance(fname, str):
        fname = (fname,)
    if get_symbol_type(symbol[0]) == "coin":
        db: str = get_db()
        table: str = get_factor_tb(timeframe)
    else:
        db: str = get_quant_db()
        table: str = get_feature_tb(timeframe)
    sql_query = f"""
            select time, fname, symbol, anyLast(value) as value from {table}
            where time>=:start and time<:end and symbol in :symbol and fname in :fname
            group by symbol, fname, time
            order by symbol, fname, time
        """
    engine = get_db_engine(db)
    with engine.connect() as connection:
        query = connection.execute(
            text(sql_query),
            {"start": start, "end": end, "symbol": symbol, "fname": fname},
        )
        results = pd.DataFrame(query.fetchall())

    results = {x: y.drop("symbol") for x, y in results.groupby("symbol")}
    results = {
        x: y.pivot(columns="fname", index="time", values="value")
        for x, y in results.items()
    }
    if is_fill:
        results = {x: fill_data(y, start, end, timeframe) for x, y in results}
    return results


def default_get_data(
    symbol: str,
    timeframe: TimeFrameType,
    start: str,
    end: str,
    preloads: int = None,
    features: list[str] = None,
):
    new_start = start
    if preloads:
        seconds = preloads * time_counts[timeframe]
        new_start = datetime.fromisoformat(start) - timedelta(seconds=seconds)
        new_start = new_start.isoformat()

    kline_features = ("open", "high", "low", "close", "volume")
    need_klines = set(features).intersection(kline_features)
    others = set(features).difference(need_klines)
    if get_symbol_type(symbol) == "coin":
        tbs = []
        if need_klines:
            kline_dfs = get_klines(new_start, end, symbol, timeframe)
            tbs.append(kline_dfs)
        if others:
            depends_df = get_factors(new_start, end, symbol, list(others), timeframe)
            tbs.append(depends_df)
        results = {x: pd.concat([tbs[0][x], tbs[1][x]]) for x in tbs[0]}
    else:
        results = get_factors(new_start, end, symbol, features, timeframe)
    return results


if __name__ == "__main__":
    # results = get_kline_unfilled_range(
    #     '2019-12-31 11:21:00',
    #     '2020-01-08 14:21:00',
    #     'BTC/USD',
    #     '1m',
    #     'test_coins',
    #     'klines_1m'
    # )
    # print(results)
    #
    # finals = get_cut_range_from_table(results, "1m", 20)
    # print(finals)

    results = get_klines(
        "2020-04-25 02:00:00", "2020-04-25 09:30:00", "BTC/USDT", "1m", is_fill=True
    )
    print(results)

    # results = get_factors("2020-11-01", "2020-11-21", "BTC/USDT", "rvi14#rvi", "1m")
    # print(results)
